{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CsVPnR8VbXE6"
      },
      "source": [
        "# Gemini API: Document Q&A with ChromaDB"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "awKO767lQIWh"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/chromadb/Vectordb_with_chroma.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" height=30/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbf4f2c17530"
      },
      "source": [
        "<!-- Princing warning Badge -->\n",
        "<table>\n",
        "  <tr>\n",
        "    <!-- Emoji -->\n",
        "    <td bgcolor=\"#f5949e\">\n",
        "      <font size=30>⚠️</font>\n",
        "    </td>\n",
        "    <!-- Text Content Cell -->\n",
        "    <td bgcolor=\"#f5949e\">\n",
        "      <h3><font color=black>This notebook requires paid tier rate limits to run properly.<br>  \n",
        "(cf. <a href=\"https://ai.google.dev/pricing#veo2\"><font color='#217bfe'>pricing</font></a> for more details).</font></h3>\n",
        "    </td>\n",
        "  </tr>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YtwZ8DZGJfUv"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use the Gemini API to create a vector database and retrieve answers to questions from the database. Moreover, you will use [ChromaDB](https://docs.trychroma.com/){:.external}, an open-source Python tool that creates embedding databases. ChromaDB allows you to:\n",
        "\n",
        "* Store embeddings as well as their metadata\n",
        "* Embed documents and queries\n",
        "* Search through the database of embeddings\n",
        "\n",
        "In this tutorial, you'll use embeddings to retrieve an answer from a database of vectors created with ChromaDB."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "akuOzK4dJl3j"
      },
      "source": [
        "## Setup\n",
        "\n",
        "First, download and install ChromaDB and the Gemini API Python library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JbXe7Oodc5dP"
      },
      "outputs": [],
      "source": [
        "%pip install -U -q \"google-genai>=1.0.0\" chromadb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jwmKt115PxK8"
      },
      "source": [
        "Then import the modules you'll use in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "muuhsDmmKdHi"
      },
      "outputs": [],
      "source": [
        "import textwrap\n",
        "import chromadb\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "from IPython.display import Markdown\n",
        "from chromadb import Documents, EmbeddingFunction, Embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FQOGMejVu-6D"
      },
      "source": [
        "## Configure your API key\n",
        "\n",
        "Before you can use the Gemini API, you must first obtain an API key. If you don't already have one, create a key with one click in Google AI Studio.\n",
        "\n",
        "<a class=\"button button-primary\" href=\"https://aistudio.google.com/app/apikey\" target=\"_blank\" rel=\"noopener noreferrer\">Get an API key</a>\n",
        "\n",
        "In Colab, add the key to the secrets manager under the \"🔑\" in the left panel. Give it the name `GEMINI_API_KEY`.\n",
        "\n",
        "Once you have the API key, pass it to the SDK. You can do this in two ways:\n",
        "\n",
        "* Put the key in the `GEMINI_API_KEY` environment variable (the SDK will automatically pick it up from there).\n",
        "* Pass the key to `genai.Client(api_key=...)`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ysayz8skEfBW"
      },
      "outputs": [],
      "source": [
        "from google import genai\n",
        "from google.colab import userdata\n",
        "\n",
        "\n",
        "GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')\n",
        "client = genai.Client(api_key=GEMINI_API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fegnGFpMS4AI"
      },
      "source": [
        "Key Point: Next, you will choose a model. Any embedding model will work for this tutorial, but for real applications it's important to choose a specific model and stick with it. The outputs of different models are not compatible with each other.\n",
        "\n",
        "**Note**: At this time, the Gemini API is [only available in certain regions](https://ai.google.dev/available_regions)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Km5d13_FS2Q_"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "models/embedding-001\n",
            "models/text-embedding-004\n",
            "models/gemini-embedding-exp-03-07\n",
            "models/gemini-embedding-exp\n",
            "models/gemini-embedding-001\n"
          ]
        }
      ],
      "source": [
        "for m in client.models.list():\n",
        "  if 'embedContent' in m.supported_actions:\n",
        "    print(m.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3XWKXoXwOGxS"
      },
      "source": [
        "### Data\n",
        "\n",
        "Here is a small set of documents you will use to create an embedding database:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "k8nsbhFJKmG-"
      },
      "outputs": [],
      "source": [
        "DOCUMENT1 = \"\"\"\n",
        "  Operating the Climate Control System  Your Googlecar has a climate control\n",
        "  system that allows you to adjust the temperature and airflow in the car.\n",
        "  To operate the climate control system, use the buttons and knobs located on\n",
        "  the center console.  Temperature: The temperature knob controls the\n",
        "  temperature inside the car. Turn the knob clockwise to increase the\n",
        "  temperature or counterclockwise to decrease the temperature.\n",
        "  Airflow: The airflow knob controls the amount of airflow inside the car.\n",
        "  Turn the knob clockwise to increase the airflow or counterclockwise to\n",
        "  decrease the airflow. Fan speed: The fan speed knob controls the speed\n",
        "  of the fan. Turn the knob clockwise to increase the fan speed or\n",
        "  counterclockwise to decrease the fan speed.\n",
        "  Mode: The mode button allows you to select the desired mode. The available\n",
        "  modes are: Auto: The car will automatically adjust the temperature and\n",
        "  airflow to maintain a comfortable level.\n",
        "  Cool: The car will blow cool air into the car.\n",
        "  Heat: The car will blow warm air into the car.\n",
        "  Defrost: The car will blow warm air onto the windshield to defrost it.\n",
        "\"\"\"\n",
        "DOCUMENT2 = \"\"\"\n",
        "  Your Googlecar has a large touchscreen display that provides access to a\n",
        "  variety of features, including navigation, entertainment, and climate\n",
        "  control. To use the touchscreen display, simply touch the desired icon.\n",
        "  For example, you can touch the \\\"Navigation\\\" icon to get directions to\n",
        "  your destination or touch the \\\"Music\\\" icon to play your favorite songs.\n",
        "\"\"\"\n",
        "DOCUMENT3 = \"\"\"\n",
        "  Shifting Gears Your Googlecar has an automatic transmission. To\n",
        "  shift gears, simply move the shift lever to the desired position.\n",
        "  Park: This position is used when you are parked. The wheels are locked\n",
        "  and the car cannot move.\n",
        "  Reverse: This position is used to back up.\n",
        "  Neutral: This position is used when you are stopped at a light or in traffic.\n",
        "  The car is not in gear and will not move unless you press the gas pedal.\n",
        "  Drive: This position is used to drive forward.\n",
        "  Low: This position is used for driving in snow or other slippery conditions.\n",
        "\"\"\"\n",
        "\n",
        "documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDzxArLeOexD"
      },
      "source": [
        "## Creating the embedding database with ChromaDB\n",
        "\n",
        "You will create a [custom function](https://docs.trychroma.com/embeddings#custom-embedding-functions){:.external} for performing embedding using the Gemini API. By inputting a set of documents into this custom function, you will receive vectors, or embeddings of the documents.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UoHhS32txd_r"
      },
      "source": [
        "### API changes to Embeddings with model gemini-embedding-001\n",
        "\n",
        "For the new embeddings model, embedding-001, there is a new task type parameter and the optional title (only valid with task_type=`RETRIEVAL_DOCUMENT`).\n",
        "\n",
        "These new parameters apply only to the newest embeddings models.The task types are:\n",
        "\n",
        "Task Type | Description\n",
        "---       | ---\n",
        "RETRIEVAL_QUERY\t| Specifies the given text is a query in a search/retrieval setting.\n",
        "RETRIEVAL_DOCUMENT | Specifies the given text is a document in a search/retrieval setting.\n",
        "SEMANTIC_SIMILARITY\t| Specifies the given text will be used for Semantic Textual Similarity (STS).\n",
        "CLASSIFICATION\t| Specifies that the embeddings will be used for classification.\n",
        "CLUSTERING\t| Specifies that the embeddings will be used for clustering."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "mF7Uu1kCQsT0"
      },
      "outputs": [],
      "source": [
        "from google.genai import types\n",
        "\n",
        "class GeminiEmbeddingFunction(EmbeddingFunction):\n",
        "  def __call__(self, input: Documents) -> Embeddings:\n",
        "    EMBEDDING_MODEL_ID = \"gemini-embedding-001\"  # @param [\"gemini-embedding-001\", \"text-embedding-004\"] {\"allow-input\": true, \"isTemplate\": true}\n",
        "    title = \"Custom query\"\n",
        "    response = client.models.embed_content(\n",
        "        model=EMBEDDING_MODEL_ID,\n",
        "        contents=input,\n",
        "        config=types.EmbedContentConfig(\n",
        "          task_type=\"retrieval_document\",\n",
        "          title=title\n",
        "        )\n",
        "    )\n",
        "\n",
        "    return response.embeddings[0].values"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrDWLyopPNBf"
      },
      "source": [
        "Now you will create the vector database. In the `create_chroma_db` function, you will instantiate a [Chroma client](https://docs.trychroma.com/getting-started). From there, you will create a collection, which is where you store your embeddings, documents, and any metadata. Note that the embedding function from above is passed as an argument to the `create_collection`.\n",
        "\n",
        "Next, you use the `add` method to add the documents to the collection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "OITXgxZlLoXU"
      },
      "outputs": [],
      "source": [
        "def create_chroma_db(documents, name):\n",
        "  chroma_client = chromadb.Client()\n",
        "  db = chroma_client.create_collection(\n",
        "      name=name,\n",
        "      embedding_function=GeminiEmbeddingFunction()\n",
        "  )\n",
        "\n",
        "  for i, d in enumerate(documents):\n",
        "    db.add(\n",
        "      documents=d,\n",
        "      ids=str(i)\n",
        "    )\n",
        "  return db"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "RJ3Fq0yzL10B"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/tmp/ipykernel_83887/781126645.py:5: DeprecationWarning: The class GeminiEmbeddingFunction does not implement __init__. This will be required in a future version.\n",
            "  embedding_function=GeminiEmbeddingFunction()\n"
          ]
        }
      ],
      "source": [
        "# Set up the DB\n",
        "db = create_chroma_db(documents, \"google-car-db\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2QbwFgfXp-fL"
      },
      "source": [
        "Confirm that the data was inserted by looking at the database:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "kQ9PHUL_l-hf"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>IDs</th>\n",
              "      <th>Documents</th>\n",
              "      <th>Embeddings</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>0</td>\n",
              "      <td>\\n  Operating the Climate Control System  Your...</td>\n",
              "      <td>[ 0.00971627 -0.00177013  0.00590323 ...  0.00...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>\\n  Your Googlecar has a large touchscreen dis...</td>\n",
              "      <td>[ 0.00388563 -0.00036349 -0.00230268 ...  0.01...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>2</td>\n",
              "      <td>\\n  Shifting Gears Your Googlecar has an autom...</td>\n",
              "      <td>[-0.00264773  0.010808   -0.00854844 ...  0.00...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  IDs                                          Documents  \\\n",
              "0   0  \\n  Operating the Climate Control System  Your...   \n",
              "1   1  \\n  Your Googlecar has a large touchscreen dis...   \n",
              "2   2  \\n  Shifting Gears Your Googlecar has an autom...   \n",
              "\n",
              "                                          Embeddings  \n",
              "0  [ 0.00971627 -0.00177013  0.00590323 ...  0.00...  \n",
              "1  [ 0.00388563 -0.00036349 -0.00230268 ...  0.01...  \n",
              "2  [-0.00264773  0.010808   -0.00854844 ...  0.00...  "
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sample_data = db.get(include=['documents', 'embeddings'])\n",
        "\n",
        "df = pd.DataFrame({\n",
        "    \"IDs\": sample_data['ids'][:3],\n",
        "    \"Documents\": sample_data['documents'][:3],\n",
        "    \"Embeddings\": [str(emb)[:50] + \"...\" for emb in sample_data['embeddings'][:3]]  # Truncate embeddings\n",
        "})\n",
        "\n",
        "df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tu5zRErgsQ8u"
      },
      "source": [
        "## Getting the relevant document\n",
        "\n",
        "`db` is a Chroma collection object. You can call `query` on it to perform a nearest neighbors search to find similar embeddings or documents.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "gQdJMbTSLtKE"
      },
      "outputs": [],
      "source": [
        "def get_relevant_passage(query, db):\n",
        "  passage = db.query(query_texts=[query], n_results=1)['documents'][0][0]\n",
        "  return passage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "nWYXXKJ6t6Hy"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "\n",
              "  Your Googlecar has a large touchscreen display that provides access to a\n",
              "  variety of features, including navigation, entertainment, and climate\n",
              "  control. To use the touchscreen display, simply touch the desired icon.\n",
              "  For example, you can touch the \"Navigation\" icon to get directions to\n",
              "  your destination or touch the \"Music\" icon to play your favorite songs.\n"
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Perform embedding search\n",
        "passage = get_relevant_passage(\"touch screen features\", db)\n",
        "Markdown(passage)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s8PNRMpOQkm5"
      },
      "source": [
        "Now that you have found the relevant passage in your set of documents, you can use it make a prompt to pass into the Gemini API."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "Qkhu4iazLy3G"
      },
      "outputs": [],
      "source": [
        "def make_prompt(query, relevant_passage):\n",
        "  escaped = relevant_passage.replace(\"'\", \"\").replace('\"', \"\").replace(\"\\n\", \" \")\n",
        "  prompt = (\"\"\"\n",
        "    You are a helpful and informative bot that answers questions using\n",
        "    text from the reference passage included below.\n",
        "    Be sure to respond in a complete sentence, being comprehensive,\n",
        "    including all relevant background information.\n",
        "    However, you are talking to a non-technical audience, so be sure to\n",
        "    break down complicated concepts and strike a friendly\n",
        "    and converstional tone. If the passage is irrelevant to the answer,\n",
        "    you may ignore it.\n",
        "    QUESTION: '{query}'\n",
        "    PASSAGE: '{relevant_passage}'\n",
        "\n",
        "    ANSWER:\n",
        "  \"\"\").format(query=query, relevant_passage=escaped)\n",
        "\n",
        "  return prompt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMEjbz4EswQ6"
      },
      "source": [
        "Pass a query to the prompt:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "b6_Y-GOymaXu"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "\n",
              "    You are a helpful and informative bot that answers questions using\n",
              "    text from the reference passage included below.\n",
              "    Be sure to respond in a complete sentence, being comprehensive,\n",
              "    including all relevant background information.\n",
              "    However, you are talking to a non-technical audience, so be sure to\n",
              "    break down complicated concepts and strike a friendly\n",
              "    and converstional tone. If the passage is irrelevant to the answer,\n",
              "    you may ignore it.\n",
              "    QUESTION: 'How do you use the touchscreen in the Google car?'\n",
              "    PASSAGE: '   Your Googlecar has a large touchscreen display that provides access to a   variety of features, including navigation, entertainment, and climate   control. To use the touchscreen display, simply touch the desired icon.   For example, you can touch the Navigation icon to get directions to   your destination or touch the Music icon to play your favorite songs. '\n",
              "\n",
              "    ANSWER:\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "query = \"How do you use the touchscreen in the Google car?\"\n",
        "prompt = make_prompt(query, passage)\n",
        "Markdown(prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VRy6yXzcPxLB"
      },
      "source": [
        "Now use the `generate_content` method to to generate a response from the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "EwfyxFM6Giy9"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "Using the touchscreen in your Google car is super easy and intuitive! You simply **touch the icon** for whatever you want to do. This large display gives you access to lots of different features, like finding your way with **navigation**, enjoying your favorite **music or entertainment**, and even controlling the **climate** inside the car. For instance, if you want directions, you would just touch the \"Navigation\" icon, or if you're in the mood for some tunes, you'd tap the \"Music\" icon to play your favorite songs."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "MODEL_ID = \"gemini-2.5-flash\" # @param [\"gemini-2.5-flash-lite\", \"gemini-2.5-flash\", \"gemini-2.5-pro\",\"gemini-3-pro-preview\"] {\"allow-input\":true, isTemplate: true}\n",
        "answer = client.models.generate_content(\n",
        "    model = MODEL_ID,\n",
        "    contents = prompt\n",
        ")\n",
        "Markdown(answer.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ThTbjAJ7eGP5"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "To learn more about how you can use the embeddings, check out the [examples](https://ai.google.dev/examples?keywords=embed) available. To learn how to use other services in the Gemini API, visit the [Python quickstart](https://ai.google.dev/gemini-api/docs/get-started/python)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Vectordb_with_chroma.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
